import argparse
from typing import Any, Dict, Optional

from airflow.patsnap.service import workflow_service
from airflow import AirflowException
from airflow.models import BaseOperator, DagModel
from airflow.utils import dates
from airflow.utils.decorators import apply_defaults

try:
    import pwd
except ImportError:
    pass


class ArgumentParserError(Exception):
    pass


class IDataArgumentParser(argparse.ArgumentParser):
    def error(self, message):
        raise ArgumentParserError(message)


class DagRunOperator(BaseOperator):

    @apply_defaults
    def __init__(
            self,
            *,
            workflow_id: Optional[int] = None,
            **kwargs: Any
    ) -> None:
        super().__init__(**kwargs)
        self._workflow_id = workflow_id
        self._execution_date = None

    def execute(self, context: Dict[str, Any]) -> None:
        execution_date = context.get("execution_date", None)
        if execution_date:
            self._execution_date = dates.to_milliseconds(execution_date)
        else:
            self._execution_date = dates.to_milliseconds(dates.sh_now())

        self.submit(execution_date=self._execution_date)

    def on_kill(self) -> None:
        self.log.debug("Kill Command is being called")

        if self._workflow_id and self._execution_date:
            try:
                dag_id = 'dag-{}'.format(self._workflow_id)
                execution_date = dates.from_milliseconds(self._execution_date)
                workflow_service.mark_dag_failed(dag_id, execution_date)
            except Exception as e:
                self.log.error("kill workflow error:{}", e)
        else:
            self.log.error("_workflow_id and _execution_date is None")

    def submit(self, execution_date, **kwargs: Any) -> None:
        if not self._workflow_id:
            raise AirflowException(
                "Cannot execute workflow_id: {}. Error message is: {}".format(
                    self._workflow_id, "_workflow_id is None"
                )
            )
        try:
            # 判断dag是否开启调度
            dag_id = 'dag-{}'.format(self._workflow_id)
            dm = DagModel.get_dagmodel(dag_id)
            if not dm.is_paused:
                raise AirflowException(
                    "Cannot execute workflow_id: {}. Error message is: {}".format(
                        self._workflow_id, "dagRun 自己dag不可开启调度"
                    )
                )
            workflow_service.execute_workflow(self._workflow_id, execution_date)
        except Exception as e:
            raise e